Skip to content

[SPARK-58736][SQL][PYTHON][CONNECT] Support incremental Python aggregators via Arrow - #57952

Open
HyukjinKwon wants to merge 12 commits into
apache:masterfrom
HyukjinKwon:SPARK-python-arrow-incremental-aggregator
Open

[SPARK-58736][SQL][PYTHON][CONNECT] Support incremental Python aggregators via Arrow#57952
HyukjinKwon wants to merge 12 commits into
apache:masterfrom
HyukjinKwon:SPARK-python-arrow-incremental-aggregator

Conversation

@HyukjinKwon

@HyukjinKwon HyukjinKwon commented Aug 12, 2026

Copy link
Copy Markdown
Member

What changes were proposed in this pull request?

Adds a Python analog of the Scala typed org.apache.spark.sql.expressions.Aggregator[IN, BUF, OUT]
with true incremental (partial) aggregation — i.e. map-side combine, not whole-group
materialization.

Users subclass a new Aggregator base class (zero / reduce / merge / finish +
bufferSchema) and wrap it with udaf(...) for use in groupBy().agg(...):

from pyspark.sql.aggregator import Aggregator, udaf
from pyspark.sql.types import StructType, StructField, DoubleType, LongType

class Mean(Aggregator):
    @property
    def bufferSchema(self):
        return StructType([StructField("sum", DoubleType()), StructField("count", LongType())])
    @property
    def outputType(self):
        return DoubleType()
    def zero(self):           return (0.0, 0)
    def reduce(self, buf, v): return (buf[0] + v[0], buf[1] + 1)
    def merge(self, a, b):    return (a[0] + b[0], a[1] + b[1])
    def finish(self, buf):    return buf[0] / buf[1] if buf[1] else None

df.groupBy("k").agg(udaf(Mean())(df.v))

Unlike grouped-agg pandas/arrow UDFs (PythonUDAF + ArrowAggregatePythonExec), which collect the
whole group and call Python once, this is planned as a two-stage aggregation:

  • a map-side PARTIAL stage folds each group's input rows into a per-group buffer via reduce;
  • the buffers are shuffled by the grouping key (as an Arrow struct column);
  • a FINAL stage merges the partial buffers via merge and produces the output via finish.

Because merge is associative/commutative, the result is independent of partition count.

Class hierarchy / trace:

  • PythonEvalType: new SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF (255) and ..._FINAL_UDF
    (256), added on both the Python (pyspark.util) and JVM (api.python.PythonEvalType) sides.
  • Catalyst: new PythonAggregate expression (an UnevaluableAggregateFunc, like PythonUDAF)
    carrying the intermediate bufferSchema.
  • Planning: SparkStrategies.Aggregation routes an all-PythonAggregate aggregate to
    PythonIncrementalAggregateExec.plan(...), which builds
    PythonIncrementalAggregatePartialExec -> (Exchange, inserted by EnsureRequirements) ->
    PythonIncrementalAggregateFinalExec. Both operators reuse ArrowPythonWithNamedArgumentRunner
    • GroupedPythonArrowInput.
  • Worker (worker.py): a PARTIAL handler that folds input batches into a buffer via reduce, and
    a FINAL handler that merges partial-buffer rows via merge then finish.
  • The buffer schema is threaded to the JVM via a new nullable bufferType on
    UserDefinedPythonFunction (an auxiliary constructor preserves the existing Py4J arity).

Spark Connect: also supported. A new optional buffer_type field on the PythonUDF proto
message carries the buffer schema to the server; the Connect client (connect/udf.py,
connect/expressions.py) serializes it and udaf dispatches on is_remote(); the server
SparkConnectPlanner threads buffer_type into UserDefinedPythonFunction and builds
PythonAggregate, after which execution reuses the same operators/worker code as classic.

SQL registration: spark.udf.register("my_agg", udaf(agg)) works in both classic and Connect,
so the aggregator is usable from SQL text (SELECT my_agg(v) FROM t GROUP BY k) — the counterpart
of Scala's spark.udf.register(name, functions.udaf(agg)).

Out of scope (planned follow-ups): DISTINCT, mixing with SQL aggregate functions in
one Aggregate, window/streaming, real disk spill (currently the map side bounds memory by
per-partition grouping; associativity makes early partial emission safe), and a typed-columnar vs.
pickled buffer performance variant.

Why are the changes needed?

PySpark has no incremental user-defined aggregator: every custom-aggregation path (grouped-agg
pandas_udf/arrow_udf, applyInPandas) materializes the whole group and invokes Python once,
with no map-side combine or partial/merge across the shuffle. This adds the missing
Aggregator-style abstraction with genuine partial aggregation, matching the Scala typed
Aggregator.

Does this PR introduce any user-facing change?

Yes — a new public API: pyspark.sql.aggregator.Aggregator and udaf(...), usable in
groupBy().agg(...) and registrable via spark.udf.register(...) for use in SQL. No existing
behavior changes.

How was this patch tested?

  • Compilation verified: sql/compile (catalyst + core + sql) builds cleanly with the new
    expression, operators, and planner routing.
  • Added python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py
    (ArrowPythonAggregatorTests): checks the incremental aggregator matches built-in avg/sum,
    a no-group case, a custom buffer, and that results are independent of partition count (exercising
    partial + merge), plus test_sql_registration (register via spark.udf.register, invoke from
    SQL text). A Connect parity suite (ArrowPythonAggregatorParityTests) runs the same mixin
    under ReusedConnectTestCase.
  • Compilation verified for both classic and Connect: sql/compile and connect/compile build
    cleanly (including proto regeneration). Full test execution runs in this PR's CI.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 4.8)

Add a Python analog of the Scala typed `Aggregator[IN, BUF, OUT]` with true
incremental (partial) aggregation. Users subclass `Aggregator`
(`zero`/`reduce`/`merge`/`finish` + `bufferSchema`) and wrap it with
`arrow_udaf(...)` for use in `groupBy().agg(...)`.

Unlike grouped-agg pandas/arrow UDFs (whole-group materialization), this is
planned as a two-stage aggregation with map-side combine: a PARTIAL stage folds
each group's input rows into a per-group Arrow buffer via `reduce`, the buffers
are shuffled by the grouping key, and a FINAL stage merges the partial buffers
via `merge` and produces the output via `finish`.

- New eval types SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL/FINAL_UDF
  (Python `PythonEvalType` and JVM `PythonEvalType`).
- New Catalyst expression `PythonAggregate` carrying the intermediate buffer
  schema (unevaluable in the JVM, like `PythonUDAF`).
- New physical operators `PythonIncrementalAggregate{Partial,Final}Exec`,
  routed in `SparkStrategies` as Partial -> Exchange -> Final; the buffer
  crosses the shuffle as an Arrow struct column.
- Worker handlers: reduce-into-buffer (partial) and merge+finish (final).
- `arrow_udaf` / `Aggregator` API under `pyspark.sql.pandas.aggregator`.

Buffer schema is threaded to the JVM via a new nullable `bufferType` on
`UserDefinedPythonFunction`. Out of scope for now (follow-ups): distinct,
mixing with SQL aggregates, window/streaming, Spark Connect, SQL registration,
and a typed-vs-pickled buffer perf variant.

Co-authored-by: Isaac
…ark Connect

Wire the incremental Python aggregator (arrow_udaf / Aggregator) through Spark
Connect so it works in remote sessions as well as classic.

- Proto: add optional `buffer_type` (DataType) to the `PythonUDF` message and
  regenerate the Python stubs.
- Connect client: `PythonUDF` expression wrapper carries `buffer_type` and
  serializes it into the proto; `UserDefinedFunction` forwards a `bufferSchema`
  attribute. `arrow_udaf` now dispatches on `is_remote()` to build the Connect
  UDF in a remote session.
- Connect server: `SparkConnectPlanner.createUserDefinedPythonFunction` threads
  `buffer_type` into `UserDefinedPythonFunction`, and `transformPythonFuncExpression`
  builds `PythonAggregate` for the incremental eval type. Execution then reuses
  the same operators/worker code as classic.
- Test: `ArrowPythonAggregatorParityTests` runs the same mixin under
  `ReusedConnectTestCase`.

Co-authored-by: Isaac
Name the factory `udaf` to mirror Scala's `functions.udaf(agg)`, and require a
supported PyArrow version up front (via require_minimum_pyarrow_version) with a
clear error, since the aggregator transfers its intermediate buffer as Arrow.

Co-authored-by: Isaac
…o 4.4.0

Relocate `aggregator.py` from `pyspark.sql.pandas` to `pyspark.sql` (import as
`pyspark.sql.aggregator`), and set the `versionadded` for `Aggregator`/`udaf`
to 4.4.0. Update the references in util.py, connect/udf.py, and the test.

Co-authored-by: Isaac
@HyukjinKwon HyukjinKwon changed the title [WIP][SQL][PYTHON] Support incremental Python aggregators via Arrow [DO-NOT-MERGE][SQL][PYTHON] Support incremental Python aggregators via Arrow Aug 12, 2026
… aggregator

Allow `spark.udf.register(name, udaf(agg))` so the incremental Python aggregator
can be invoked from SQL text (`SELECT my_agg(v) FROM t GROUP BY k`), matching
Scala's `spark.udf.register(name, functions.udaf(agg))`.

- Classic and Connect `UDFRegistration.register` accept
  SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF and thread the buffer schema
  through (classic `register` reconstructs the UDF and would otherwise drop it;
  Connect passes it via `SparkConnectClient.register_udf` -> the PythonUDF proto).
- `udaf` sets `bufferSchema` on the returned wrapper too, so it survives
  registration. The Connect server already builds `PythonAggregate` in
  `handleRegisterUserDefinedFunction` via the shared `createUserDefinedPythonFunction`.
- Test: `test_sql_registration` in the shared mixin (runs classic + Connect).

Co-authored-by: Isaac
@HyukjinKwon

Copy link
Copy Markdown
Member Author

cc @zhengruifeng @cloud-fan @Yicong-Huang Seems like this way it can do the actual partial aggregation.

Comment thread python/pyspark/worker.py
# profiling is not supported for UDF
return grouped_func, None, ser, ser

if eval_type == PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why can't this reuse PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is a different iterator? it is for element-iterator inside a row, not a row-iterator

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah my take is that PythonEvalType.* decides internal computation type

@HyukjinKwon

Copy link
Copy Markdown
Member Author
GROUPED-AGG pandas UDF            single stage · groupByKey
───────────────────────────────────────────────────────────
     P1[a b a]     P2[b a b]     P3[a a b]     3 partitions
         └──────────────┼──────────────┘
                        ▼
     ═══ SHUFFLE ═══ all 9 raw rows move ═══
                ┌───────┴───────┐
                ▼               ▼
       ┌────────────────┐  ┌────────────────┐
       │ key a          │  │ key b          │
       │ a a a a a      │  │ b b b b        │  ← whole group,
       │ → udf(Series)  │  │ → udf(Series)  │    one worker
       └────────────────┘  └────────────────┘
                a → r               b → r


INCREMENTAL Aggregator (udaf)     two stages
───────────────────────────────────────────────────────────
     P1[a b a]     P2[b a b]     P3[a a b]
       │ reduce       │ reduce       │ reduce  ┐
       ▼              ▼              ▼
    [Σa Σb]        [Σa Σb]        [Σa Σb]      ┘ (map-side combine)
         └──────────────┼──────────────┘
                        ▼
     ═══ SHUFFLE ═══ only 6 buffers move ═══
                ┌───────┴───────┐
                ▼               ▼
       ┌────────────────┐  ┌────────────────┐
       │ key a          │  │ key b
       │ Σa Σa Σa       │  │ Σb Σb Σb       │  ← only a few
       │ → merge → fin  │  │ → merge → fin
       └────────────────┘  └────────────────┘
                a → r               b → r

- Use PySparkNotImplementedError instead of a raw NotImplementedError in
  Aggregator.__call__ (PySpark custom-errors linter).
- Import have_pyarrow / pyarrow_requirement_message from pyspark.testing.utils
  (not sqlutils), which was causing the aggregator test modules to fail at import.

Co-authored-by: Isaac
@HyukjinKwon HyukjinKwon changed the title [DO-NOT-MERGE][SQL][PYTHON] Support incremental Python aggregators via Arrow [SPARK-58736][SQL][PYTHON][CONNECT] Support incremental Python aggregators via Arrow Aug 12, 2026
@HyukjinKwon
HyukjinKwon marked this pull request as ready for review August 12, 2026 10:40

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2 blocking, 0 non-blocking, 0 nits.
The two-stage integration is coherent, but the worker lifecycle still has two blocking semantic gaps in incremental memory use and empty-input aggregation.

Design / architecture (1)

  • Blocking: python/pyspark/worker.py:2234: Stream partial input batches into the aggregation buffers instead of retaining and concatenating the complete group first. -- see inline

Correctness (1)

  • Blocking: sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonIncrementalAggregateExec.scala:126: Emit the identity buffer for empty global input so finish(zero) produces the required single aggregate row. -- see inline

Verification

Traced udaf through classic and Connect UDF construction, Catalyst planning, both physical stages, and the Python worker handlers. Confirmed that the partial handler calls list(group) before reduce, and that the physical operator returns an empty iterator before the identity buffer can be emitted for empty global input. Tests were not run as part of this review.

Comment thread python/pyspark/worker.py Outdated
…obal aggregation

Two blocking review items:

- PARTIAL/FINAL worker handlers now stream Arrow batches and fold them one at a
  time into the per-aggregator buffers, instead of `list(group)` + concatenating
  the whole group first. Map-side peak memory is bounded by a single batch (plus
  the buffers), not the whole group -- the point of the incremental API.
- A global (no-grouping) aggregation over empty input now returns the identity
  row `finish(zero)` instead of no row. GroupedPythonArrowInput cannot transmit
  an empty group, so the FINAL stage (which runs on a single AllTuples partition)
  injects one all-null buffer row; the worker skips null partial buffers and so
  merges nothing, yielding `finish(zero)`. Added a focused test
  (`df.limit(0).agg(udaf(...))`).

Co-authored-by: Isaac
- invalidPandasUDFPlacementError now also names incremental PythonAggregate
  functions (not just grouped-agg PythonUDAF) when Python aggregate UDFs are
  mixed with other aggregate functions in one Aggregate.
- Add a test with two incremental aggregators (different buffer schemas) over
  the same input, covering multi-UDF partial/final planning and execution.

Co-authored-by: Isaac
…message

- Define ArrowGroupedAggIncremental{Partial,Final}UDFType Literal aliases and
  import them under TYPE_CHECKING so the eval-type annotations resolve (F821).
- Use the standard `from pyspark.testing import main` test footer instead of
  `import *` (F403 / RUF100); reformat with ruff.
- Update the INVALID_UDF_EVAL_TYPE expected message in test_pandas_grouped_map
  to include SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF now that the
  incremental aggregator is registerable via spark.udf.register.

Co-authored-by: Isaac
Reformat with ruff 0.14.0 to match the CI-pinned version (files were
previously formatted with an older local ruff).

Co-authored-by: Isaac
Silence mypy's [assignment] error on the classic UserDefinedFunction import
in the is_remote() dispatch, matching the Connect/classic dispatch pattern.

Co-authored-by: Isaac
@HyukjinKwon

Copy link
Copy Markdown
Member Author

Code review (self, head 1d6c9391)

Verdict: approve with minor cleanups. No correctness or design blockers. The two-stage design is coherent, both earlier blocking comments (stream batches into buffers; emit finish(zero) for empty global input) are genuinely fixed, and Connect parity + SQL registration are wired through. Findings below are all minor.

Findings

1. FINAL worker handler builds output arrays without an explicit Arrow typerobustness

In python/pyspark/worker.py, the FINAL handler does:

result_arrays = [pa.array([r]) for r in results]

whereas the PARTIAL handler correctly passes type=return_schema.field(i).type. Relying purely on enforce_schema to coerce an inferred type is fragile for non-trivial outputTypes (decimal, timestamp, nested struct) or an all-None column. Suggest aligning FINAL with PARTIAL:

result_arrays = [pa.array([r], type=return_schema.field(i).type) for i, r in enumerate(results)]

2. Loop-invariant field_names recomputation in both worker handlersminor perf

field_names = [f.name for f in agg.bufferSchema.fields] is recomputed per (group × batch × aggregator) in the FINAL handler and per (group × aggregator) in the PARTIAL handler, though it depends only on the fixed i-th aggregator. Precompute a field_names_by_udf list once where grouped_func is defined and index by i.

3. Unqualified Scaladoc linkdoc nit

In sql/catalyst/.../expressions/PythonUDF.scala, [[PythonIncrementalAggregateExec]] cannot resolve from sql/catalyst (the class lives in sql/core). Fully-qualify it as [[org.apache.spark.sql.execution.python.PythonIncrementalAggregateExec]] (as the reverse-direction reference already does) or use plain text.

4. aggregator.py docstring polishoptional

  • The reduce example (buffer[0] + v, buffer[1] + 1) raises on a null input value; consider showing null handling so the example isn't copied as a fragile pattern.
  • udaf's Raises: lists only PySparkImportError, but it also raises PySparkTypeError for a non-Aggregator arg or a non-StructType bufferSchema.

Test portfolio

Strong coverage: builtin mean/sum equivalence, no-group, empty-global input, custom buffer, multiple aggregators, partition-count independence, SQL registration, and a Connect parity mirror. Two gaps worth a follow-up: no test with a non-trivial output type (decimal/timestamp) — which would exercise finding #1 — and no test of reduce receiving a null input value.

…d names, doc/link fixes

- worker.py FINAL handler: type each output array explicitly via
  return_schema.field(i).type (mirroring PARTIAL) instead of relying on Arrow
  type inference + enforce_schema, so a non-trivial outputType (e.g. decimal)
  or an all-None column is robust.
- worker.py PARTIAL and FINAL: hoist the loop-invariant per-aggregator
  bufferSchema field-name lists (field_names_by_udf) out of the group/batch
  loops.
- PythonUDF.scala: fully-qualify the [[...PythonIncrementalAggregateExec]]
  Scaladoc link, which could not resolve from sql/catalyst (the class is in
  sql/core).
- aggregator.py: show null-input handling in the reduce example and document
  that udaf raises PySparkTypeError for a bad agg / bufferSchema.
- Tests: add a DecimalType-output aggregator test (exercises explicit output
  typing across the shuffle) and a null-input test; make the example Mean skip
  nulls to match SQL avg semantics.

Co-authored-by: Isaac
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants